import torch
import torch.nn as nn
import torch.nn.functional as F


class GeneratorGRU(nn.Module):
    def __init__(self, embed_dim, ff1_out, hidden_dim, out_dim):
        super(GeneratorGRU, self).__init__()

        self.input_ff = nn.Linear(embed_dim, ff1_out)
        self.lstm = nn.GRU(ff1_out, hidden_dim, num_layers=2)
        self.output_ff = nn.Linear(hidden_dim, out_dim)

    def forward(self, lyrics):
        """
        Define forward pass
        1. Pass input (50 dim) through FF1 layer to explode the dimension
        2. Pass the entire sequence through
        :return:
        """
        # print("Input size {}".format(lyrics.shape))
        # Reshaping input is not required.
        # pytorch automatically applys the linear layer to only the last dimension!

        # (100,10,32)
        # print("Shape of input lyrics is: {}".format(lyrics.shape))
        # (100,10,32)
        # print("Shape of noise is: {}".format(noise.shape))
        concat_lyrics = torch.cat((lyrics, torch.FloatTensor(lyrics.size()).uniform_()), 2)
        # print('concat_lyrics size:{}'.format(concat_lyrics.shape))
        # (100,10,64)
        # print("Concat Shape: {}".format(concat_lyrics.shape))
        out_input_ff = self.input_ff(concat_lyrics)
        # print('out_input_ff size:{}'.format(out_input_ff.shape))
        # print(out_input_ff.shape)#(100,10,400)
        out1 = F.relu(out_input_ff)
        # print("Output size of first layer {}".format(out1.shape))
        # print(out1.shape)
        # (100,10,400)
        # lstm_out, _ = self.lstm(out1)
        # The input to LSTM needs to be reshaped.
        lstm_out, _ = self.lstm(out1.view(out1.shape[1], out1.shape[0], -1))
        # print('lstm_out size:{}'.format(lstm_out.shape))
        # print(lstm_out.shape)
        # (10,100,400)

        tag = self.output_ff(lstm_out.view(lstm_out.shape[1], lstm_out.shape[0], -1))
        # print('tag size:{}'.format(tag.shape))
        # print(tag.shape)
        # (100,10,3)

        # print(tag)
        return tag

    # concat_lyrics
    # size: torch.Size([100, 10, 64])
    # out_input_ff
    # size: torch.Size([100, 10, 32])
    # lstm_out
    # size: torch.Size([10, 100, 400])
    # tag
    # size: torch.Size([100, 10, 3])
